#pragma once

#include <assert.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <vector>

#ifdef _WIN32
#include <windows.h>

#else
#include <sys/mman.h>
#include <unistd.h>
#include <pthread.h>
#include <malloc.h>

#endif

// tcmalloc中用三颗基数树, 优化map搜索时的锁紧竞争
// TCMalloc中使用基数树优化了对内存页号的管理，减少了在多线程情况下的锁竞争。

// 存储页号需要多少位, BITS = 32/64 - PAGE_SHIFT
// 在32位系统中，BITS为(32 - PAGE_SHIFT)，PAGE_SHIFT一般是12，即2^12=4096字节

// Single-level array
// 单级数组实现的基数树
template <int BITS>
class TCMalloc_PageMap1
{
private:
    static const int LENGTH = 1 << BITS; // 数组长度为2的BITS次方
    void** array_; // 存放指针的数组

public:
    typedef uintptr_t Number; // 定义页号类型为uintptr_t

    explicit TCMalloc_PageMap1() {
        array_ = new void* [LENGTH]; // 使用new分配内存
        memset(array_, 0, sizeof(void*) * LENGTH); // 初始化数组，全部置为0
    }

    ~TCMalloc_PageMap1() {
        delete[] array_; // 释放内存
    }

    // 返回键为k的当前值。如果未设置或k超出范围，返回NULL。
    void* get(Number k) const {
        if ((k >> BITS) > 0) { // 如果k超出范围（大于等于1<<BITS）
            return NULL;
        }
        return array_[k]; // 返回数组中索引为k的元素
    }

    // 设置键为k的值为v。
    void set(Number k, void* v) {
        array_[k] = v; // 设置数组中索引为k的元素为v
    }

    void erase(Number k) const {
        if ((k >> BITS) > 0) { // 如果k超出范围（大于等于1<<BITS）
            assert(false);
        }
        array_[k] = nullptr;
    }
};

// 两级基数树
template <int BITS>
class TCMalloc_PageMap2
{
private:
    static const int ROOT_BITS = 5; // 根节点位数，决定根节点数量（32个）
    static const int ROOT_LENGTH = 1 << ROOT_BITS; // 根节点数组长度
    static const int LEAF_BITS = BITS - ROOT_BITS; // 叶子节点位数
    static const int LEAF_LENGTH = 1 << LEAF_BITS; // 每个叶子节点数组长度

    // 叶子节点结构体
    struct Leaf
    {
        void* values[LEAF_LENGTH]; // 叶子节点存放值的数组
    };

    Leaf* root_[ROOT_LENGTH]; // 根节点指针数组，指向32个子节点（Leaf*类型）

public:
    typedef uintptr_t Number; // 定义页号类型为uintptr_t

    explicit TCMalloc_PageMap2() {
        memset(root_, 0, sizeof(root_)); // 初始化根节点数组，全部置为NULL
        PreallocateMoreMemory(); // 预分配更多内存
    }

    ~TCMalloc_PageMap2() {
        for (int i = 0; i < ROOT_LENGTH; ++i) {
            delete root_[i];
        }
    }

    // 返回键为k的当前值。如果未设置或k超出范围，返回NULL。
    void* get(Number k) const {
        const Number i1 = k >> LEAF_BITS; // 计算根节点索引
        const Number i2 = k & (LEAF_LENGTH - 1); // 计算叶子节点索引
        if ((k >> BITS) > 0 || root_[i1] == NULL) { // 如果k超出范围或根节点为空
            return NULL;
        }
        return root_[i1]->values[i2]; // 返回叶子节点中索引为i2的值
    }

    // 设置键为k的值为v。
    void set(Number k, void* v) {
        const Number i1 = k >> LEAF_BITS; // 计算根节点索引
        const Number i2 = k & (LEAF_LENGTH - 1); // 计算叶子节点索引
        assert(i1 < ROOT_LENGTH); // 断言根节点索引小于ROOT_LENGTH
        root_[i1]->values[i2] = v; // 设置叶子节点中索引为i2的值为v
    }

    void erase(Number k) const {
        const Number i1 = k >> LEAF_BITS; // 计算根节点索引
        const Number i2 = k & (LEAF_LENGTH - 1); // 计算叶子节点索引
        if ((k >> BITS) > 0 || root_[i1] == NULL) { // 如果k超出范围或根节点为空
            assert(false);
        }
        root_[i1]->values[i2] = nullptr;
    }

    // 确保从start开始的n个页面被分配
    bool Ensure(Number start, size_t n) {
        for (Number key = start; key <= start + n - 1;) {
            const Number i1 = key >> LEAF_BITS; // 计算根节点索引
            if (i1 >= ROOT_LENGTH) return false; // 检查是否溢出

            // 如果需要，创建第二级节点
            if (root_[i1] == NULL) {
                Leaf* leaf = new Leaf(); // 使用new分配内存
                memset(leaf, 0, sizeof(*leaf)); // 初始化Leaf节点，全部置为0
                root_[i1] = leaf; // 将Leaf节点赋给根节点数组对应位置
            }
            // 将key前进到此叶子节点所覆盖的位置之后
            key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
        }
        return true;
    }

    // 预分配更多内存，确保能够管理所有可能的页面
    void PreallocateMoreMemory() {
        Ensure(0, 1 << BITS); // 确保分配所有可能页面的内存
    }
};


// Three-level radix tree
// 三级基数树
template <int BITS>
class TCMalloc_PageMap3
{
private:
    // How many bits should we consume at each interior level
    // 每个内部级别应该使用多少位
    static const int INTERIOR_BITS = (BITS + 2) / 3; // Round-up
    static const int INTERIOR_LENGTH = 1 << INTERIOR_BITS; // 中间节点数组长度

    // How many bits should we consume at leaf level
    // 每个叶子级别应该使用多少位
    static const int LEAF_BITS = BITS - 2 * INTERIOR_BITS; // 叶子节点位数
    static const int LEAF_LENGTH = 1 << LEAF_BITS; // 叶子节点数组长度

    // Interior node
    // 中间节点结构体
    struct Node
    {
        Node* ptrs[INTERIOR_LENGTH]; // 指向中间节点数组
    };

    // Leaf node
    // 叶子节点结构体
    struct Leaf
    {
        void* values[LEAF_LENGTH]; // 叶子节点存放值的数组
    };

    Node* root_; // Root of radix tree，基数树的根节点


    // 记录已分配节点的位置
    std::vector<Node*> allocated_nodes_;
    std::vector<Leaf*> allocated_leaves_;

    // Allocate a new Node
    // 分配一个新的Node节点
    Node* NewNode() {
        Node* result = new Node(); // 使用new分配内存
        memset(result, 0, sizeof(*result)); // 初始化Node节点，全部置为0

        allocated_nodes_.push_back(result); // 记录已分配的节点

        return result;
    }

public:
    typedef uintptr_t Number; // 定义页号类型为uintptr_t

    // Constructor with allocator function pointer
    // 构造函数，接受一个allocator函数指针作为参数
    explicit TCMalloc_PageMap3() {
        root_ = NewNode(); // 分配根节点
        PreallocateMoreMemory();
    }

    // 析构函数
    ~TCMalloc_PageMap3() {
        for (auto node : allocated_nodes_) {
            delete node; // 释放所有已分配的Node节点
        }
        for (auto leaf : allocated_leaves_) {
            delete leaf; // 释放所有已分配的Leaf节点
        }
    }

    // Get the value associated with key k
    // 返回与键k关联的值
    void* get(Number k) const {
        const Number i1 = k >> (LEAF_BITS + INTERIOR_BITS); // 计算第一级索引
        const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1); // 计算第二级索引
        const Number i3 = k & (LEAF_LENGTH - 1); // 计算第三级索引

        // Check if key k is out of range or nodes are not allocated
        // 检查键k是否超出范围或节点未分配
        if ((k >> BITS) > 0 ||
            root_->ptrs[i1] == NULL || root_->ptrs[i1]->ptrs[i2] == NULL) {
            return NULL;
        }

        // Return the value stored in the leaf node
        // 返回叶子节点中存储的值
        return reinterpret_cast<Leaf*>(root_->ptrs[i1]->ptrs[i2])->values[i3];
    }

    // Set the value v for key k
    // 设置键k的值为v
    void set(Number k, void* v) {
        assert((k >> BITS) == 0); // 断言确保键k在有效范围内

        const Number i1 = k >> (LEAF_BITS + INTERIOR_BITS); // 计算第一级索引
        const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1); // 计算第二级索引
        const Number i3 = k & (LEAF_LENGTH - 1); // 计算第三级索引

        // Ensure that necessary nodes are allocated
        // 确保必要的节点已分配
        if (root_->ptrs[i1] == NULL || root_->ptrs[i1]->ptrs[i2] == NULL) {
            Ensure(k, 1); // 分配必要节点
        }

        // Set the value in the leaf node
        // 在叶子节点中设置值v
        reinterpret_cast<Leaf*>(root_->ptrs[i1]->ptrs[i2])->values[i3] = v;
    }

    void erase(Number k) const {
        const Number i1 = k >> (LEAF_BITS + INTERIOR_BITS); // 计算第一级索引
        const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1); // 计算第二级索引
        const Number i3 = k & (LEAF_LENGTH - 1); // 计算第三级索引

        // Check if key k is out of range or nodes are not allocated
        // 检查键k是否超出范围或节点未分配
        if ((k >> BITS) > 0 || root_->ptrs[i1] == NULL || root_->ptrs[i1]->ptrs[i2] == NULL) {
            assert(false);
        }
        reinterpret_cast<Leaf*>(root_->ptrs[i1]->ptrs[i2])->values[i3] = nullptr;
    }

    // Ensure that n pages starting from start are allocated
    // 确保从start开始的n个页面被分配
    bool Ensure(Number start, size_t n) {
        for (Number key = start; key <= start + n - 1;) {
            const Number i1 = key >> (LEAF_BITS + INTERIOR_BITS); // 计算第一级索引
            const Number i2 = (key >> LEAF_BITS) & (INTERIOR_LENGTH - 1); // 计算第二级索引

            // Check for overflow
            // 检查是否溢出
            if (i1 >= INTERIOR_LENGTH || i2 >= INTERIOR_LENGTH)
                return false;

            // Make 2nd level node if necessary
            // 如有必要，创建第二级节点
            if (root_->ptrs[i1] == NULL) {
                Node* n = NewNode(); // 分配新的Node节点
                if (n == NULL) return false;
                root_->ptrs[i1] = n; // 将新节点赋给根节点数组对应位置
            }

            // Make leaf node if necessary
            // 如有必要，创建叶子节点
            if (root_->ptrs[i1]->ptrs[i2] == NULL) {
                Leaf* leaf = new Leaf; // 分配内存
                if (leaf == NULL) return false;

                memset(leaf, 0, sizeof(*leaf)); // 初始化Leaf节点，全部置为0
                root_->ptrs[i1]->ptrs[i2] = reinterpret_cast<Node*>(leaf); // 将Leaf节点赋给根节点数组对应位置

                allocated_leaves_.push_back(leaf); // 记录已分配的叶子节点
            }

            // Advance key past whatever is covered by this leaf node
            // 将key前进到此叶子节点所覆盖的位置之后
            key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
        }
        return true;
    }

    // Preallocate more memory if needed
    // 如果需要，预分配更多内存
    void PreallocateMoreMemory() {
        // This function can be implemented to preallocate memory as required
        // 可以根据需要实现此函数来预分配内存


    }
};